import numpy as np
import os
import json
from Embedding import Embedding

class Search:
    
    def __init__(self):
        self.__globalKnowledge__ = []
        self.emb = Embedding()
        # self.embedding()
        self.__load__()
        
        
    def embedding(self):
        for file in os.listdir(os.path.join("RawData", "Global")):
            with open(os.path.join("RawData", "Global", file), "r", encoding="utf-8") as f:
                content = [i + "。" for i in f.read().split("。")]
                embeddingResult = [[sentence, self.emb.embedding(sentence)] for sentence in content]
                self.__globalKnowledge__.extend(embeddingResult)
                with open(os.path.join("repo", "global", file[:-len(".txt")] + ".json"), "w", encoding="utf-8") as output:
                    json.dump(embeddingResult, output, ensure_ascii=False)
                    
    def __load__(self):
        for file in os.listdir(os.path.join("repo", "global")):
            content = json.load(open(os.path.join("repo", "global", file), "r", encoding="utf-8"))
            self.__globalKnowledge__.extend(content)
            
    def search(self, query) -> list:
        query = self.emb.embedding(query)
        consineSimilarity = []
        for index in range(len(self.__globalKnowledge__)):
            consineSimilarity.append([index, self.__cosine_similarity__(query, self.__globalKnowledge__[index][1])])
        consineSimilarity.sort(key=lambda x: x[1], reverse=True)
        return [self.__globalKnowledge__[index][0] for index, sentence in consineSimilarity[:4]]
    
    @staticmethod
    def __cosine_similarity__(vector1, vector2):
        # 将列表转换为numpy数组
        vec1 = np.array(vector1)
        vec2 = np.array(vector2)
        
        # 计算向量的点积
        dot_product = np.dot(vec1, vec2)
        
        # 计算向量的模长
        norm_vec1 = np.linalg.norm(vec1)
        norm_vec2 = np.linalg.norm(vec2)
        
        # 计算余弦相似度
        similarity = dot_product / (norm_vec1 * norm_vec2)
        
        return similarity
